import pandas as pd
from pprint import pprint


class DataOperation:
    def __init__(self, file_name):
        # 获取文件格式，打开文件
        a, b = file_name.rsplit('.')
        if str(b) == 'csv' or str(b) == 'tsv' or str(b) == 'txt':
            self.data = pd.read_csv(file_name)
        elif str(b) == 'xls' or str(b) == 'xlsx':
            self.data = pd.read_excel(file_name)
        else:
            print('请输入正确文件格式')

    def read_n_row(self, n):
        # 获取前n行
        pprint(self.data.head(n))

    def read_index(self):
        # 获取行索引
        pprint(self.data.index)

    def read_shape(self):
        # 获取数据形状，多少行，多少列
        print(self.data.shape)

    def read_columns(self):
        # 查看字段名
        print(self.data.columns)

    def view_by_label(self, *args):
        # 通过行标签和列标签查看数据
        try:
            pprint(self.data.loc[args])
        except KeyError:
            print('输入的标签有误')

    def view_by_index(self, *args):
        # 通过行索引列索引查看数据
        try:
            pprint(self.data.iloc[args])
        except IndexError:
            print('输入的索引有误')

    def read_data_return_list(self):
        # 读取数据返回列表
        pprint(self.data.iloc[self.data.index.values].values.tolist())

    def read_data_return_dict(self):
        # 读取数据返回字典
        # list2 = []
        # for i in self.data.index.values:
        #     list2.append(self.data.loc[i].to_dict())
        # pprint(list2)
        list3 = [self.data.loc[i].to_dict() for i in self.data.index.values]
        pprint(list3)


if __name__ == '__main__':
    data = DataOperation('111.xlsx')
    # data1 = DataOperation('123.txt')
    # data.read_n_row(1)
    # data.read_index()
    # data.read_shape()
    # data.read_columns()
    # data.view_by_label([1,2],['姓名'])
    # data.view_by_index([1, 2], [1])
    data.read_data_return_list()
    # data.read_data_return_dict()
